import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import random
from collections import deque, namedtuple
import heapq

# 通用组件
class ReplayBuffer:
    """
    标准经验回放缓冲区
    """
    def __init__(self, capacity):
        self.buffer = deque(maxlen=capacity)
    
    def add(self, state, action, reward, next_state, done):
        experience = (state, action, reward, next_state, done)
        self.buffer.append(experience)
    
    def sample(self, batch_size):
        batch = random.sample(self.buffer, batch_size)
        states, actions, rewards, next_states, dones = zip(*batch)
        return np.array(states), np.array(actions), np.array(rewards), np.array(next_states), np.array(dones)
    
    def __len__(self):
        return len(self.buffer)

class PrioritizedReplayBuffer:
    """
    优先级经验回放缓冲区
    """
    def __init__(self, capacity, alpha=0.6, beta=0.4, beta_increment=0.001):
        self.capacity = capacity
        self.buffer = []
        self.pos = 0
        self.priorities = np.zeros((capacity,), dtype=np.float32)
        self.alpha = alpha
        self.beta = beta
        self.beta_increment = beta_increment
        self.experience = namedtuple("Experience", field_names=["state", "action", "reward", "next_state", "done"])
    
    def add(self, state, action, reward, next_state, done):
        max_priority = self.priorities.max() if self.buffer else 1.0
        
        if len(self.buffer) < self.capacity:
            self.buffer.append(self.experience(state, action, reward, next_state, done))
        else:
            self.buffer[self.pos] = self.experience(state, action, reward, next_state, done)
        
        self.priorities[self.pos] = max_priority
        self.pos = (self.pos + 1) % self.capacity
    
    def sample(self, batch_size):
        if len(self.buffer) < self.capacity:
            priorities = self.priorities[:len(self.buffer)]
        else:
            priorities = self.priorities
        
        probs = priorities ** self.alpha
        probs /= probs.sum()
        
        indices = np.random.choice(len(self.buffer), batch_size, p=probs)
        experiences = [self.buffer[idx] for idx in indices]
        
        self.beta = np.min([1.0, self.beta + self.beta_increment])
        
        total = len(self.buffer)
        weights = (total * probs[indices]) ** (-self.beta)
        weights /= weights.max()
        
        states = np.array([exp.state for exp in experiences])
        actions = np.array([exp.action for exp in experiences])
        rewards = np.array([exp.reward for exp in experiences])
        next_states = np.array([exp.next_state for exp in experiences])
        dones = np.array([exp.done for exp in experiences])
        
        return states, actions, rewards, next_states, dones, indices, weights
    
    def update_priorities(self, indices, priorities):
        for idx, priority in zip(indices, priorities):
            self.priorities[idx] = priority
    
    def __len__(self):
        return len(self.buffer)

# 基础网络模型
class DQNNetwork(nn.Module):
    """
    基础DQN网络模型
    """
    def __init__(self, state_dim, action_dim, hidden_dim=64):
        super(DQNNetwork, self).__init__()
        self.fc1 = nn.Linear(state_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
        self.fc3 = nn.Linear(hidden_dim, action_dim)
    
    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

class DuelingDQNNetwork(nn.Module):
    """
    Dueling DQN网络模型
    """
    def __init__(self, state_dim, action_dim, hidden_dim=64):
        super(DuelingDQNNetwork, self).__init__()
        self.feature_layer = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.ReLU()
        )
        
        # 价值流
        self.value_stream = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 1)
        )
        
        # 优势流
        self.advantage_stream = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, action_dim)
        )
    
    def forward(self, x):
        features = self.feature_layer(x)
        value = self.value_stream(features)
        advantage = self.advantage_stream(features)
        
        # 结合价值和优势
        q_values = value + (advantage - advantage.mean(dim=1, keepdim=True))
        return q_values

class ActorCriticNetwork(nn.Module):
    """
    Actor-Critic网络模型
    """
    def __init__(self, state_dim, action_dim, hidden_dim=64):
        super(ActorCriticNetwork, self).__init__()
        
        # 共享特征层
        self.feature_layer = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.ReLU()
        )
        
        # Actor网络 (策略)
        self.actor = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, action_dim),
            nn.Softmax(dim=-1)
        )
        
        # Critic网络 (价值)
        self.critic = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 1)
        )
    
    def forward(self, x):
        features = self.feature_layer(x)
        policy_dist = self.actor(features)
        value = self.critic(features)
        return policy_dist, value

# DQN及其变体
class DQNAgent:
    """
    基础DQN智能体
    """
    def __init__(self, state_dim, action_dim, hidden_dim=64, lr=1e-3, gamma=0.99, 
                 epsilon=1.0, epsilon_min=0.01, epsilon_decay=0.995, batch_size=64, 
                 buffer_capacity=10000, target_update_freq=10):
        self.state_dim = state_dim
        self.action_dim = action_dim
        self.gamma = gamma
        self.epsilon = epsilon
        self.epsilon_min = epsilon_min
        self.epsilon_decay = epsilon_decay
        self.batch_size = batch_size
        self.target_update_freq = target_update_freq
        
        # 创建Q网络
        self.q_network = DQNNetwork(state_dim, action_dim, hidden_dim)
        self.target_network = DQNNetwork(state_dim, action_dim, hidden_dim)
        self.target_network.load_state_dict(self.q_network.state_dict())
        self.target_network.eval()
        
        # 优化器
        self.optimizer = torch.optim.Adam(self.q_network.parameters(), lr=lr)
        
        # 经验回放缓冲区
        self.replay_buffer = ReplayBuffer(buffer_capacity)
    
    def select_action(self, state, evaluate=False):
        if not evaluate and np.random.rand() <= self.epsilon:
            return random.randrange(self.action_dim)
        
        state_tensor = torch.FloatTensor(state).unsqueeze(0)
        with torch.no_grad():
            q_values = self.q_network(state_tensor)
        return np.argmax(q_values.cpu().data.numpy())
    
    def update(self):
        if len(self.replay_buffer) < self.batch_size:
            return 0
        
        # 从经验回放缓冲区采样
        states, actions, rewards, next_states, dones = self.replay_buffer.sample(self.batch_size)
        
        # 转换为张量
        states_tensor = torch.FloatTensor(states)
        actions_tensor = torch.LongTensor(actions).unsqueeze(1)
        rewards_tensor = torch.FloatTensor(rewards)
        next_states_tensor = torch.FloatTensor(next_states)
        dones_tensor = torch.FloatTensor(dones)
        
        # 计算当前Q值
        current_q = self.q_network(states_tensor).gather(1, actions_tensor).squeeze(1)
        
        # 计算目标Q值
        with torch.no_grad():
            max_next_q = self.target_network(next_states_tensor).max(1)[0]
            target_q = rewards_tensor + (1 - dones_tensor) * self.gamma * max_next_q
        
        # 计算损失
        loss = F.mse_loss(current_q, target_q)
        
        # 反向传播
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()
        
        # 衰减探索率
        self.epsilon = max(self.epsilon_min, self.epsilon * self.epsilon_decay)
        
        return loss.item()
    
    def update_target_network(self):
        self.target_network.load_state_dict(self.q_network.state_dict())
    
    def save_model(self, path):
        torch.save(self.q_network.state_dict(), path)
    
    def load_model(self, path):
        self.q_network.load_state_dict(torch.load(path))
        self.target_network.load_state_dict(self.q_network.state_dict())

class DoubleDQNAgent(DQNAgent):
    """
    Double DQN智能体
    """
    def update(self):
        if len(self.replay_buffer) < self.batch_size:
            return 0
        
        # 从经验回放缓冲区采样
        states, actions, rewards, next_states, dones = self.replay_buffer.sample(self.batch_size)
        
        # 转换为张量
        states_tensor = torch.FloatTensor(states)
        actions_tensor = torch.LongTensor(actions).unsqueeze(1)
        rewards_tensor = torch.FloatTensor(rewards)
        next_states_tensor = torch.FloatTensor(next_states)
        dones_tensor = torch.FloatTensor(dones)
        
        # 计算当前Q值
        current_q = self.q_network(states_tensor).gather(1, actions_tensor).squeeze(1)
        
        # Double DQN: 使用当前网络选择动作，使用目标网络评估
        with torch.no_grad():
            # 当前网络选择最佳动作
            best_actions = self.q_network(next_states_tensor).argmax(1, keepdim=True)
            # 目标网络评估这些动作
            max_next_q = self.target_network(next_states_tensor).gather(1, best_actions).squeeze(1)
            target_q = rewards_tensor + (1 - dones_tensor) * self.gamma * max_next_q
        
        # 计算损失
        loss = F.mse_loss(current_q, target_q)
        
        # 反向传播
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()
        
        # 衰减探索率
        self.epsilon = max(self.epsilon_min, self.epsilon * self.epsilon_decay)
        
        return loss.item()

class DuelingDQNAgent(DQNAgent):
    """
    Dueling DQN智能体
    """
    def __init__(self, state_dim, action_dim, hidden_dim=64, lr=1e-3, gamma=0.99, 
                 epsilon=1.0, epsilon_min=0.01, epsilon_decay=0.995, batch_size=64, 
                 buffer_capacity=10000, target_update_freq=10):
        # 调用父类初始化
        super(DuelingDQNAgent, self).__init__(state_dim, action_dim, hidden_dim, lr, gamma, 
                                             epsilon, epsilon_min, epsilon_decay, batch_size, 
                                             buffer_capacity, target_update_freq)
        
        # 替换为Dueling DQN网络
        self.q_network = DuelingDQNNetwork(state_dim, action_dim, hidden_dim)
        self.target_network = DuelingDQNNetwork(state_dim, action_dim, hidden_dim)
        self.target_network.load_state_dict(self.q_network.state_dict())
        self.target_network.eval()
        
        # 重新初始化优化器
        self.optimizer = torch.optim.Adam(self.q_network.parameters(), lr=lr)

class PERDQNAgent(DQNAgent):
    """
    带优先级经验回放的DQN智能体
    """
    def __init__(self, state_dim, action_dim, hidden_dim=64, lr=1e-3, gamma=0.99, 
                 epsilon=1.0, epsilon_min=0.01, epsilon_decay=0.995, batch_size=64, 
                 buffer_capacity=10000, target_update_freq=10, alpha=0.6, beta=0.4):
        # 调用父类初始化
        super(PERDQNAgent, self).__init__(state_dim, action_dim, hidden_dim, lr, gamma, 
                                         epsilon, epsilon_min, epsilon_decay, batch_size, 
                                         buffer_capacity, target_update_freq)
        
        # 替换为优先级经验回放缓冲区
        self.replay_buffer = PrioritizedReplayBuffer(buffer_capacity, alpha, beta)
    
    def update(self):
        if len(self.replay_buffer) < self.batch_size:
            return 0
        
        # 从优先级经验回放缓冲区采样
        states, actions, rewards, next_states, dones, indices, weights = self.replay_buffer.sample(self.batch_size)
        
        # 转换为张量
        states_tensor = torch.FloatTensor(states)
        actions_tensor = torch.LongTensor(actions).unsqueeze(1)
        rewards_tensor = torch.FloatTensor(rewards)
        next_states_tensor = torch.FloatTensor(next_states)
        dones_tensor = torch.FloatTensor(dones)
        weights_tensor = torch.FloatTensor(weights)
        
        # 计算当前Q值
        current_q = self.q_network(states_tensor).gather(1, actions_tensor).squeeze(1)
        
        # 计算目标Q值
        with torch.no_grad():
            max_next_q = self.target_network(next_states_tensor).max(1)[0]
            target_q = rewards_tensor + (1 - dones_tensor) * self.gamma * max_next_q
        
        # 计算TD误差用于更新优先级
        td_errors = torch.abs(current_q - target_q).detach().numpy()
        
        # 使用权重计算加权损失
        loss = (weights_tensor * F.mse_loss(current_q, target_q, reduction='none')).mean()
        
        # 反向传播
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()
        
        # 更新优先级
        self.replay_buffer.update_priorities(indices, td_errors + 1e-6)  # 避免零优先级
        
        # 衰减探索率
        self.epsilon = max(self.epsilon_min, self.epsilon * self.epsilon_decay)
        
        return loss.item()

# Policy Gradient方法
class A2CAgent:
    """
    Advantage Actor-Critic (A2C)智能体
    """
    def __init__(self, state_dim, action_dim, hidden_dim=64, lr=3e-4, gamma=0.99, entropy_coef=0.01):
        self.state_dim = state_dim
        self.action_dim = action_dim
        self.gamma = gamma
        self.entropy_coef = entropy_coef
        
        # 创建Actor-Critic网络
        self.network = ActorCriticNetwork(state_dim, action_dim, hidden_dim)
        
        # 优化器
        self.optimizer = torch.optim.Adam(self.network.parameters(), lr=lr)
    
    def select_action(self, state, evaluate=False):
        state_tensor = torch.FloatTensor(state).unsqueeze(0)
        
        with torch.no_grad():
            policy_dist, _ = self.network(state_tensor)
            probs = policy_dist.cpu().numpy()[0]
        
        if evaluate:
            # 评估模式选择概率最高的动作
            return np.argmax(probs)
        else:
            # 训练模式根据概率分布采样动作
            return np.random.choice(self.action_dim, p=probs)
    
    def update(self, states, actions, rewards, next_states, dones):
        # 转换为张量
        states_tensor = torch.FloatTensor(states)
        actions_tensor = torch.LongTensor(actions)
        rewards_tensor = torch.FloatTensor(rewards)
        next_states_tensor = torch.FloatTensor(next_states)
        dones_tensor = torch.FloatTensor(dones)
        
        # 前向传播
        policy_dists, values = self.network(states_tensor)
        
        # 计算TD目标
        with torch.no_grad():
            _, next_values = self.network(next_states_tensor)
            td_targets = rewards_tensor + (1 - dones_tensor) * self.gamma * next_values.squeeze()
        
        # 计算优势
        advantages = td_targets - values.squeeze()
        
        # 计算策略损失
        log_probs = torch.log(policy_dists.gather(1, actions_tensor.unsqueeze(1))).squeeze()
        policy_loss = -(log_probs * advantages.detach()).mean()
        
        # 计算价值损失
        value_loss = F.mse_loss(values.squeeze(), td_targets.detach())
        
        # 计算熵正则化项
        entropy = -(policy_dists * torch.log(policy_dists + 1e-8)).sum(1).mean()
        
        # 总损失
        total_loss = policy_loss + 0.5 * value_loss - self.entropy_coef * entropy
        
        # 反向传播
        self.optimizer.zero_grad()
        total_loss.backward()
        self.optimizer.step()
        
        return total_loss.item(), policy_loss.item(), value_loss.item(), entropy.item()
    
    def save_model(self, path):
        torch.save(self.network.state_dict(), path)
    
    def load_model(self, path):
        self.network.load_state_dict(torch.load(path))

# 训练和测试函数
def train_agent(env, agent, episodes=100, max_steps=1000, target_update_freq=10, update_every=1):
    """
    通用智能体训练函数
    """
    rewards_history = []
    losses_history = []
    
    for episode in range(episodes):
        state = env.reset()
        total_reward = 0
        done = False
        episode_losses = []
        steps = 0
        
        # A2C的轨迹存储
        if isinstance(agent, A2CAgent):
            states = []
            actions = []
            rewards = []
            next_states = []
            dones = []
        
        while not done and steps < max_steps:
            # 选择动作
            action = agent.select_action(state)
            
            # 执行动作
            next_state, reward, done, _ = env.step(action)
            
            if isinstance(agent, A2CAgent):
                # A2C收集轨迹
                states.append(state)
                actions.append(action)
                rewards.append(reward)
                next_states.append(next_state)
                dones.append(done)
                
                # 每一步更新A2C
                if steps % update_every == 0 and steps > 0:
                    loss_info = agent.update(
                        states[-update_every:], 
                        actions[-update_every:], 
                        rewards[-update_every:], 
                        next_states[-update_every:], 
                        dones[-update_every:]
                    )
                    episode_losses.append(loss_info[0])  # 记录总损失
            else:
                # DQN系列收集经验并更新
                agent.replay_buffer.add(state, action, reward, next_state, done)
                loss = agent.update()
                if loss > 0:
                    episode_losses.append(loss)
            
            # 更新状态和总奖励
            state = next_state
            total_reward += reward
            steps += 1
        
        # DQN系列更新目标网络
        if not isinstance(agent, A2CAgent) and episode % target_update_freq == 0:
            agent.update_target_network()
        
        rewards_history.append(total_reward)
        if episode_losses:
            losses_history.append(np.mean(episode_losses))
        
        # 打印训练进度
        if episode % 10 == 0:
            avg_reward = np.mean(rewards_history[-10:]) if len(rewards_history) >= 10 else total_reward
            avg_loss = np.mean(losses_history[-10:]) if len(losses_history) >= 10 else 0
            print(f"Episode {episode}, Total Reward: {total_reward:.2f}, Avg Reward: {avg_reward:.2f}")
            if isinstance(agent, DQNAgent):
                print(f"  Epsilon: {agent.epsilon:.3f}")
            if avg_loss > 0:
                print(f"  Avg Loss: {avg_loss:.4f}")
    
    return rewards_history

def test_agent(env, agent, episodes=10):
    """
    通用智能体测试函数
    """
    rewards_history = []
    
    # 保存原始参数
    if hasattr(agent, 'epsilon'):
        original_epsilon = agent.epsilon
        agent.epsilon = 0.0  # 测试时不探索
    
    for episode in range(episodes):
        state = env.reset()
        total_reward = 0
        done = False
        
        while not done:
            action = agent.select_action(state, evaluate=True)
            next_state, reward, done, _ = env.step(action)
            state = next_state
            total_reward += reward
        
        rewards_history.append(total_reward)
        print(f"Test Episode {episode}, Total Reward: {total_reward:.2f}")
    
    # 恢复原始参数
    if hasattr(agent, 'epsilon'):
        agent.epsilon = original_epsilon
    
    avg_reward = np.mean(rewards_history)
    print(f"\n测试完成! 平均奖励: {avg_reward:.2f}")
    
    return rewards_history, avg_reward

# 创建智能体的工厂函数
def create_agent(agent_type, state_dim, action_dim, **kwargs):
    """
    创建指定类型的智能体
    
    Args:
        agent_type: 智能体类型 ('dqn', 'double_dqn', 'dueling_dqn', 'per_dqn', 'a2c')
        state_dim: 状态维度
        action_dim: 动作维度
        **kwargs: 其他参数
        
    Returns:
        创建的智能体实例
    """
    agents = {
        'dqn': DQNAgent,
        'double_dqn': DoubleDQNAgent,
        'dueling_dqn': DuelingDQNAgent,
        'per_dqn': PERDQNAgent,
        'a2c': A2CAgent
    }
    
    if agent_type.lower() not in agents:
        raise ValueError(f"不支持的智能体类型: {agent_type}. 支持的类型: {list(agents.keys())}")
    
    return agents[agent_type.lower()](state_dim, action_dim, **kwargs)

# 运行示例的主函数
def run_rl_example(agent_type='dqn', env_name='CartPole-v1', episodes=100):
    """
    运行强化学习算法示例
    
    Args:
        agent_type: 智能体类型
        env_name: 环境名称
        episodes: 训练轮数
        
    Returns:
        训练奖励历史
    """
    try:
        import gym
        
        # 创建环境
        env = gym.make(env_name)
        
        # 获取状态和动作维度
        state_dim = env.observation_space.shape[0]
        action_dim = env.action_space.n
        
        print(f"\n开始训练 {agent_type.upper()} 智能体在 {env_name} 环境中...")
        print(f"状态维度: {state_dim}, 动作维度: {action_dim}")
        
        # 设置特定算法的参数
        if agent_type.lower() == 'a2c':
            kwargs = {'lr': 3e-4, 'entropy_coef': 0.01}
        else:
            kwargs = {'lr': 1e-3, 'batch_size': 64}
        
        # 创建智能体
        agent = create_agent(agent_type, state_dim, action_dim, **kwargs)
        
        # 训练智能体
        train_rewards = train_agent(env, agent, episodes)
        
        print("\n开始测试智能体...")
        # 测试智能体
        test_rewards, avg_reward = test_agent(env, agent)
        
        # 关闭环境
        env.close()
        
        return train_rewards, test_rewards, avg_reward
        
    except ImportError:
        print("请先安装gym库: pip install gym")
        return None, None, None
    except Exception as e:
        print(f"运行示例时出错: {e}")
        return None, None, None

# 主函数示例
if __name__ == "__main__":
    run_rl_example()